use cudarc::driver::{
    CudaContext, CudaModule, CudaSlice, CudaStream, DriverError, LaunchConfig, PushKernelArg,
};
use glam::{Mat4, Quat, Vec3, Vec4, vec3, vec4};
use glium::{
    Display,
    winit::{
        event::{ElementState, KeyEvent, WindowEvent},
        event_loop::EventLoop,
        keyboard::{self, NamedKey},
    },
};
use glutin::surface::WindowSurface;
use std::{
    path::{Path, PathBuf},
    sync::Arc,
};
use use_cudarc::{
    framebuffer_renderer::{CudaRendererOps, FrameBufferRenderer},
    model::{GpuModel, Model},
    tgaimage::{ColorFormat, TGAImage},
    window_app::App,
};

const Z_BUFFER_MIN: f32 = -1e6;

struct ModelRenderer {
    frame_buffer_size: (usize, usize),
    frame_buffer: Vec<u8>,
    frame_buffer_gpu: CudaSlice<u8>,
    z_buffer: CudaSlice<f32>,

    model_view: CudaSlice<f32>,
    perspective: CudaSlice<f32>,
    viewport: CudaSlice<f32>,

    cuda_stream: Arc<CudaStream>,
    module: Arc<CudaModule>,
}

impl ModelRenderer {
    fn new<P: AsRef<Path>>(
        frame_buffer_size: (usize, usize),
        shader: P,
    ) -> Result<Self, DriverError> {
        let cuda_ctx = CudaContext::new(0)?;
        let cuda_stream = cuda_ctx.default_stream();

        let frame_buffer = vec![[0, 0, 0, 255]; frame_buffer_size.0 * frame_buffer_size.1].concat();
        let frame_buffer_gpu = cuda_stream.memcpy_stod(&frame_buffer)?;
        let z_buffer =
            cuda_stream.memcpy_stod(&vec![0.; frame_buffer_size.0 * frame_buffer_size.1])?;

        let ptx = cudarc::nvrtc::compile_ptx(std::fs::read_to_string(shader).unwrap()).unwrap();
        let module = cuda_ctx.load_module(ptx)?;

        let model_view = cuda_stream.memcpy_stod(&vec![0f32; 4 * 4])?;
        let perspective = cuda_stream.memcpy_stod(&vec![0f32; 4 * 4])?;
        let viewport = cuda_stream.memcpy_stod(&vec![0f32; 4 * 4])?;

        Ok(Self {
            frame_buffer_size,
            frame_buffer,
            frame_buffer_gpu,
            z_buffer,
            cuda_stream,
            module,
            model_view,
            perspective,
            viewport,
        })
    }

    fn render_model(&mut self, model: &GpuModel) -> Result<(), DriverError> {
        let kernel = self.module.load_function("render_model_kernel")?;

        let face_num = (model.facet_vrt.len() / 3) as u32;
        let (width, height) = (
            self.frame_buffer_size.0 as u32,
            self.frame_buffer_size.1 as u32,
        );

        let mut builder = self.cuda_stream.launch_builder(&kernel);
        builder.arg(&model.verts);
        builder.arg(&model.facet_vrt);
        builder.arg(&face_num);
        builder.arg(&mut self.frame_buffer_gpu);
        builder.arg(&mut self.z_buffer);
        builder.arg(&width);
        builder.arg(&height);
        builder.arg(&self.model_view);
        builder.arg(&self.perspective);
        builder.arg(&self.viewport);
        #[cfg(feature = "perf_test")]
        let t0 = std::time::Instant::now();
        #[cfg(feature = "perf_test")]
        {
            self.cuda_stream.synchronize().unwrap();
        }
        unsafe { builder.launch(LaunchConfig::for_num_elems(face_num)) }?;
        #[cfg(feature = "perf_test")]
        {
            self.cuda_stream.synchronize().unwrap();
            let time_in_kernel = (std::time::Instant::now() - t0).as_secs_f32() * 1000.;
            println!(
                "execute kernel used time = {time_in_kernel} ms, expect fps = {}",
                1000. / time_in_kernel
            )
        }

        self.cuda_stream
            .memcpy_dtoh(&self.frame_buffer_gpu, &mut self.frame_buffer)
    }

    fn z_buffer(&self) -> Vec<u8> {
        let z_buffer = self.cuda_stream.memcpy_dtov(&self.z_buffer).unwrap();
        let (min, max) = (
            z_buffer
                .iter()
                .filter(|x| **x > Z_BUFFER_MIN)
                .fold(f32::INFINITY, |a, &b| a.min(b)),
            z_buffer.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)),
        );
        z_buffer
            .iter()
            .map(|x| {
                let gray = ((x - min) / (max - min) * 255.) as u8;
                vec![gray, gray, gray, 255]
            })
            .collect::<Vec<_>>()
            .concat()
    }

    fn reset_framebuffer(&mut self) {
        self.cuda_stream
            .memset_zeros(&mut self.frame_buffer_gpu)
            .unwrap();
        self.cuda_stream
            .memcpy_htod(
                &vec![Z_BUFFER_MIN; self.frame_buffer_size.0 * self.frame_buffer_size.1],
                &mut self.z_buffer,
            )
            .unwrap();
    }

    fn export_framebuffer<P: AsRef<Path>>(&self, framebuffer_path: P, z_buffer_path: P) {
        let (width, height) = self.frame_buffer_size;

        TGAImage::new(width, height, ColorFormat::RGBA)
            .set_data(&self.frame_buffer)
            .write_tga_file(framebuffer_path, true)
            .unwrap();
        TGAImage::new(width, height, ColorFormat::RGBA)
            .set_data(&self.z_buffer())
            .write_tga_file(z_buffer_path, true)
            .unwrap();
    }

    fn set_model_view(&mut self, model_view: &Mat4) {
        self.cuda_stream
            .memcpy_htod(model_view.as_ref(), &mut self.model_view)
            .unwrap();
    }

    fn set_perspective(&mut self, perspective: &Mat4) {
        self.cuda_stream
            .memcpy_htod(perspective.as_ref(), &mut self.perspective)
            .unwrap();
    }

    fn set_viewport(&mut self, viewport: &Mat4) {
        self.cuda_stream
            .memcpy_htod(viewport.as_ref(), &mut self.viewport)
            .unwrap();
    }
}

struct SceneRenderer {
    model_renderer: ModelRenderer,
    models: Vec<GpuModel>,
    camera: Vec3,
}

impl SceneRenderer {
    fn gen_eye(&self) -> Vec3 {
        let pos0 = vec3(-1., 0., 2.);
        let rot = Quat::from_rotation_x(self.camera.x) * Quat::from_rotation_y(self.camera.y);
        (rot * pos0) * self.camera.z / 8.
    }
}

impl CudaRendererOps for SceneRenderer {
    fn handle_window_event(&mut self, event: &WindowEvent, _window: &glium::winit::window::Window) {
        if let WindowEvent::KeyboardInput {
            event:
                KeyEvent {
                    state: ElementState::Pressed,
                    logical_key: keyboard::Key::Named(key),
                    ..
                },
            ..
        } = event
        {
            match key {
                NamedKey::ArrowUp => self.camera.x += 0.01,
                NamedKey::ArrowDown => self.camera.x -= 0.01,
                NamedKey::ArrowRight => self.camera.y += 0.01,
                NamedKey::ArrowLeft => self.camera.y -= 0.01,
                NamedKey::PageUp => self.camera.z *= 1.05,
                NamedKey::PageDown => self.camera.z /= 1.05,
                _ => {}
            }
        }
    }

    fn gen_framebuffer(&mut self, frame_index: usize, fb_size: (usize, usize)) -> Vec<u8> {
        assert!(self.model_renderer.frame_buffer_size == fb_size);

        self.model_renderer.reset_framebuffer();

        let (width, height) = self.model_renderer.frame_buffer_size;

        let model_view = Mat4::look_at_rh(self.gen_eye(), Vec3::ZERO, Vec3::Y).transpose();
        self.model_renderer.set_model_view(&model_view);

        let perspective =
            Mat4::from_cols(Vec4::X, Vec4::Y, vec4(0., 0., 1., -1.), Vec4::W).transpose();
        self.model_renderer.set_perspective(&perspective);

        let view_port_gen = |x: f32, y: f32, width: f32, height: f32| {
            Mat4::from_cols(
                vec4(width / 2., 0., 0., x + width / 2.),
                vec4(0., height / 2., 0., y + height / 2.),
                Vec4::Z,
                Vec4::W,
            )
        };
        self.model_renderer.set_viewport(&view_port_gen(
            (width as f32) / 16.,
            (height as f32) / 16.,
            (width as f32) * 7. / 8.,
            (height as f32) * 7. / 8.,
        ));

        for model in &self.models {
            self.model_renderer.render_model(model).unwrap()
        }

        if frame_index == 0 || frame_index == 100 {
            self.model_renderer.export_framebuffer(
                format!("frame{frame_index}.tga"),
                format!("frame{frame_index}_z_buffer.tga"),
            );
        }

        self.model_renderer.frame_buffer.clone()
    }
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let model_dir = &args[1];
    assert!(std::fs::exists(model_dir).unwrap());

    let event_loop = EventLoop::new().unwrap();

    let bf_size = (1080, 1080);
    let model_renderer = ModelRenderer::new(
        bf_size,
        PathBuf::from(file!()).parent().unwrap().join("kernel.cu"),
    )
    .unwrap();

    let mut models = vec![];
    for entry in std::fs::read_dir(model_dir).unwrap() {
        let entry = entry.unwrap();
        let path = entry.path();

        if path.is_file() && path.extension().map(|ext| ext == "obj").unwrap_or(false) {
            models.push(
                GpuModel::from_model(
                    &model_renderer.cuda_stream,
                    &Model::from_file(&path).unwrap(),
                )
                .unwrap(),
            );
        }
    }

    let mut app = App::new(
        Box::new(
            move |display: &Display<WindowSurface>| -> FrameBufferRenderer<SceneRenderer> {
                FrameBufferRenderer::<SceneRenderer>::new(
                    display,
                    bf_size,
                    SceneRenderer {
                        model_renderer,
                        models,
                        camera: Vec3::Z,
                    },
                )
            },
        ),
        "CUDA RENDER",
        bf_size,
    );
    event_loop.run_app(&mut app).unwrap()
}
